#include "pcagent_utils.h"
#include <tchar.h>
#include <atlpath.h>
#include "../base/constants.h"
#include "../base/reg_key.h"
#include "../base/debug.h"
#include "../base/error.h"
#include "../base/system.h"
#include "../base/scope_guard.h"
#include "../base/service_utils.h"
#include "../common/config_manager.h"

namespace pcagent_utils
{
	HRESULT RedirectHKCR(bool is_machine)
	{
		RegKey classes_key;
		HRESULT hr = classes_key.Open(is_machine ? HKEY_LOCAL_MACHINE : HKEY_CURRENT_USER, 
			_T("Software\\Classes"), KEY_ALL_ACCESS);
		if (FAILED(hr))
		{
			ASSERT(FALSE, (_T("RedirectHKCR - key.Open(%d) fail %d"), is_machine, hr));
			return hr;
		}

		LONG result = ::RegOverridePredefKey(HKEY_CLASSES_ROOT, classes_key.Key());
		if (result != ERROR_SUCCESS)
		{
			ASSERT(FALSE, (_T("RedirectHKCR - RegOverridePredefKey fail %d"), result));
			return HRESULT_FROM_WIN32(result);
		}
		return S_OK;
	}

	HRESULT RemoveRedirectHKCR()
	{
		LONG result = ::RegOverridePredefKey(HKEY_CLASSES_ROOT, NULL);
		if (result != ERROR_SUCCESS)
		{
			ASSERT(FALSE, (_T("RedirectHKCR - RegOverridePredefKey fail %d"), result));
			return HRESULT_FROM_WIN32(result);
		}
		return S_OK;
	}

	HRESULT RegisterTypeLibForUser(ITypeLib* lib,
		OLECHAR* path,
		OLECHAR* help_dir)
	{
			CORE_LOG(L3, (_T("[RegisterTypeLibForUser]")));
			ASSERT1(lib);
			ASSERT1(path);
			// help_dir can be NULL.

			const TCHAR* library_name = _T("oleaut32.dll");
			scoped_library module(static_cast<HINSTANCE>(::LoadLibrary(library_name)));
			if (!module) {
				HRESULT hr = HRESULTFromLastError();
				CORE_LOG(LEVEL_ERROR,
					(_T("[LoadLibrary failed][%s][0x%08x]"), library_name, hr));
				return hr;
			}

			// RegisterTypeLibForUser function from oleaut32.dll.
			typedef HRESULT(__stdcall *PF)(ITypeLib*, OLECHAR*, OLECHAR*);

			const char* function_name = "RegisterTypeLibForUser";
			PF fp = reinterpret_cast<PF>(::GetProcAddress(get(module), function_name));
			if (!fp) {
				HRESULT hr = HRESULTFromLastError();
				CORE_LOG(LEVEL_ERROR,
					(_T("[GetProcAddress failed][%s][0x%08x]"),
					function_name, library_name, hr));
				return hr;
			}

			CORE_LOG(L3, (_T("[Calling RegisterTypelibForUser in oleaut]")));
			HRESULT hr = fp(lib, path, help_dir);
			if (FAILED(hr)) {
				CORE_LOG(LEVEL_ERROR, (_T("[regtypelib_for_user failed][0x%08x]"), hr));
				return hr;
			}

			return S_OK;
	}

	HRESULT UnRegisterTypeLibForUser(REFGUID lib_id,
		WORD major_ver_num,
		WORD minor_ver_num,
		LCID lcid,
		SYSKIND syskind)
	{
			CORE_LOG(L3, (_T("[UnRegisterTypeLibForUser]")));

			const TCHAR* library_name = _T("oleaut32.dll");
			scoped_library module(static_cast<HINSTANCE>(::LoadLibrary(library_name)));
			if (!module) {
				HRESULT hr = HRESULTFromLastError();
				CORE_LOG(LEVEL_ERROR,
					(_T("[LoadLibrary failed][%s][0x%08x]"), library_name, hr));
				return hr;
			}

			// UnRegisterTypeLibForUser function from oleaut32.dll.
			typedef HRESULT (__stdcall *PF)(REFGUID, WORD, WORD, LCID, SYSKIND);

			const char* function_name = "UnRegisterTypeLibForUser";
			PF fp = reinterpret_cast<PF>(::GetProcAddress(get(module), function_name));
			if (!fp) {
				HRESULT hr = HRESULTFromLastError();
				CORE_LOG(LEVEL_ERROR,
					(_T("[GetProcAddress failed][%s][0x%08x]"),
					function_name, library_name, hr));
				return hr;
			}

			CORE_LOG(L3, (_T("[Calling UnRegisterTypeLibForUser in oleaut]")));
			HRESULT hr = fp(lib_id, major_ver_num, minor_ver_num, lcid, syskind);
			if (FAILED(hr)) {
				CORE_LOG(LEVEL_ERROR, (_T("[unregtypelib_for_user failed][0x%08x]"), hr));
				return hr;
			}

			return S_OK;
	}

	HRESULT RegisterTypeLib(bool is_admin,
		const CComBSTR& path,
		ITypeLib* type_lib)
	{
		// Typelib registration.
		CORE_LOG(L3, (_T("[Registering TypeLib]")));
		HRESULT hr = S_OK;
		if (!is_admin &&
			SUCCEEDED(pcagent_utils::RegisterTypeLibForUser(type_lib, path, NULL))) 
				return S_OK;

		// For Admin cases, we use ::RegisterTypeLib().
		// For platforms where ::RegisterTypeLibForUser is not available, we register
		// with ::RegisterTypeLib, and rely on HKCR=>HKCU redirection.
		hr = ::RegisterTypeLib(type_lib, path, NULL);
		ASSERT(SUCCEEDED(hr), (_T("[TypeLib registration failed][0x%08x]"), hr));
		return hr;
	}

	HRESULT UnRegisterTypeLib(bool is_admin, const CComBSTR&, ITypeLib* type_lib)
	{
		// Typelib unregistration.
		CORE_LOG(L3, (_T("[Unregistering Typelib]")));
		TLIBATTR* tlib_attr = NULL;
		HRESULT hr = type_lib->GetLibAttr(&tlib_attr);
		ASSERT(SUCCEEDED(hr), (_T("[GetLibAttr failed][0x%08x]"), hr));
		if (FAILED(hr)) 
			return hr;
		
		ON_SCOPE_EXIT_OBJ(*type_lib, &ITypeLib::ReleaseTLibAttr, tlib_attr);

		if (!is_admin &&
			SUCCEEDED(pcagent_utils::UnRegisterTypeLibForUser(
			tlib_attr->guid,
			tlib_attr->wMajorVerNum,
			tlib_attr->wMinorVerNum,
			tlib_attr->lcid,
			tlib_attr->syskind))) {
				return S_OK;
		}

		// For Admin cases, we use ::UnRegisterTypeLib().
		// For platforms where ::UnRegisterTypeLibForUser is not available, we
		// unregister with ::UnRegisterTypeLib, and rely on HKCR=>HKCU redirection.
		hr = ::UnRegisterTypeLib(tlib_attr->guid,
			tlib_attr->wMajorVerNum,
			tlib_attr->wMinorVerNum,
			tlib_attr->lcid,
			tlib_attr->syskind);

		// We assert before the check for TYPE_E_REGISTRYACCESS below because we want
		// to catch the case where we're trying to unregister more than once because
		// that would be a bug.
		ASSERT(SUCCEEDED(hr),
			(_T("[UnRegisterTypeLib failed.  ")
			_T("This is likely a multiple unregister bug.][0x%08x]"), hr));

		// If you try to unregister a type library that's already unregistered,
		// it will return with this failure, which is OK.
		if (hr == TYPE_E_REGISTRYACCESS) {
			hr = S_OK;
		}

		return hr;
	}

	HRESULT RegisterOrUnregisterModule(bool is_machine,
		bool register_server,
		RegisterOrUnregisterFunction registrar,
		void* data)
	{
		ASSERT1(registrar);

		// ATL by default registers the control to HKCR and we want to register
		// either in HKLM, or in HKCU, depending on whether we are laying down
		// the system PCAgent, or the user PCAgent.
		// We solve this for the user PCAgent case by:
		// * Having the RGS file take a HKROOT parameter that translates to either
		//   HKLM or HKCU.
		// * Redirecting HKCR to HKCU\software\classes, for a user installation, to
		//   cover Proxy registration.
		// For the machine case, we still redirect HKCR to HKLM\\Software\\Classes,
		// to ensure that Proxy registration happens in HKLM.
		HRESULT hr = RedirectHKCR(is_machine);
		ASSERT1(SUCCEEDED(hr));
		if (FAILED(hr)) return hr;

		// We need to stop redirecting at the end of this function.
		ON_SCOPE_EXIT(RemoveRedirectHKCR);

		hr = (*registrar)(data, register_server);
		if (FAILED(hr))
		{
			CORE_LOG(LW, (_T("[RegisterOrUnregisterModule failed][%d][0x%08x]"), register_server, hr));
			ASSERT1(!register_server);
		}
		return hr;
	}

	HRESULT RegisterOrUnregisterModuleWithTypelib(
		bool is_machine,
		bool register_server,
		RegisterOrUnregisterFunction registrar,
		void* data) 
	{
		ASSERT1(registrar);

		// By default, ATL registers the control to HKCR and we want to register
		// either in HKLM, or in HKCU, depending on whether we are laying down
		// the machine PCAgent, or the user PCAgent.
		// We solve this for the user PCAgent case by:
		// * Having the RGS file take a HKROOT parameter that translates to either
		//   HKLM or HKCU.
		// * Redirecting HKCR to HKCU\software\classes, for a user installation, to
		//   cover AppId and TypeLib registration
		// * All the above makes ATL work correctly for 2K/XP. However on Win2K3
		//   and Vista, redirection does not work by itself, because in these
		//   platforms, RegisterTypeLib writes explicitly to HKLM\Software\Classes.
		//   We need to specifically call the new RegisterTypeLibForUser() API.
		//   So, we do that as well.
		// For the machine case, we still redirect HKCR to HKLM\\Software\\Classes,
		// because otherwise RegisterTypeLib ends up overwriting HKCU if the key
		// already exists in HKCU.
		HRESULT hr = RedirectHKCR(is_machine);
		ASSERT1(SUCCEEDED(hr));
		if (FAILED(hr)) {
			return hr;
		}
		// We need to stop redirecting at the end of this function.
		ON_SCOPE_EXIT(RemoveRedirectHKCR);

		// load the type library.
		CComPtr<ITypeLib> type_lib;
		CComBSTR path;
		hr = ::AtlLoadTypeLib(_AtlBaseModule.GetModuleInstance(), NULL, &path,
			&type_lib);
		if (FAILED(hr)) {
			ASSERT(false, (_T("[AtlLoadTypeLib failed][0x%08x]"), hr));
			return hr;
		}

		if (register_server) {
			hr = (*registrar)(data, register_server);
			if (FAILED(hr)) {
				ASSERT(false, (_T("[Module registration failed][0x%08x]"), hr));
				return hr;
			}

			return RegisterTypeLib(is_machine, path, type_lib);
		} else {
			hr = UnRegisterTypeLib(is_machine, path, type_lib);
			if (FAILED(hr)) {
				ASSERT(false, (_T("[UnRegisterTypeLib failed][0x%08x]"), hr));
				return hr;
			}

			return (*registrar)(data, register_server);
		}
	}

	// Creates a unique name of the form "{prefix}1c9b3d6baf90df3" and stores it in
	// the registry under HKLM/HKCU\Longkey\PCAgent\value_name. Subsequent
	// invocations of GetCurrentTaskName() will return this new value.
	HRESULT CreateAndSetVersionedNameInRegistry(bool is_machine, 
		const TCHAR* prefix, const TCHAR* value_name)
	{
		ASSERT1(prefix && *prefix);
		ASSERT1(value_name && *value_name);

		CString name(ServiceInstall::GenerateServiceName(prefix));
		CORE_LOG(L3, (_T("Versioned name[%s][%s][%s]"), prefix, value_name, name));

		const TCHAR* key_name = is_machine ? MACHINE_REG_UPDATE : USER_REG_UPDATE;
		return RegKey::SetValue(key_name, value_name, name);
	}

	CString BuildLongkeyPCAgentExeDir(bool is_machine) {
		ConfigManager& cm = *ConfigManager::Instance();
		return is_machine ? cm.GetMachinePCAgentInstallDir() :
			cm.GetUserPCAgentInstallDir();
	}

	CString BuildLongkeyPCAgentExePath(bool is_machine) {
		CORE_LOG(L3, (_T("[BuildLongkeyPCAgentExePath][%d]"), is_machine));

		CPath full_file_path(BuildLongkeyPCAgentExeDir(is_machine));
		VERIFY1(full_file_path.Append(kPCAgentShellFileName));

		return full_file_path;
	}

	HRESULT StartPCAgentWithArgs(bool is_machine,
		const TCHAR* args,
		HANDLE* process) {
			CORE_LOG(L3, (_T("[StartPCAgentWithArgs][%d][%s]"),
				is_machine, args ? args : _T("")));

			CString exe_path = BuildLongkeyPCAgentExePath(is_machine);

			CORE_LOG(L3, (_T("[command line][%s][%s]"), exe_path, args ? args : _T("")));

			HRESULT hr = System::ShellExecuteProcess(exe_path, args, NULL, process);
			if (FAILED(hr)) {
				CORE_LOG(LE, (_T("[can't start process][%s][0x%08x]"), exe_path, hr));
				return hr;
			}
			return S_OK;
	}

	// TODO: This method's name is much more specific than what it does. Can
	// we just copy the code to scheduled task and service code and eliminate it?
	// Reads the current value under {HKLM|HKCU}\Google\Update\value_name. Returns
	// default_val if value_name does not exist.
	CString GetCurrentVersionedName(bool is_machine,
		const TCHAR* value_name,
		const TCHAR* default_val)
	{
		CORE_LOG(L3, (_T("[ConfigManager::GetCurrentVersionedName]")));
		ASSERT1(value_name && *value_name);
		ASSERT1(default_val && *default_val);

		const TCHAR* key_name = is_machine ? MACHINE_REG_UPDATE : USER_REG_UPDATE;
		CString name;
		HRESULT hr(RegKey::GetValue(key_name, value_name, &name));
		if (FAILED(hr)) {
			CORE_LOG(L4, (_T("[GetValue failed][%s][0x%x][Using default name][%s]"),
				value_name, hr, default_val));
			name = default_val;
		}

		CORE_LOG(L3, (_T("[Versioned Name][%s]"), name));
		return name;
	}

}